#!/bin/bash

# PPO强化学习训练自动重启脚本
# 适用于处理OOM和其他崩溃情况 - 无清理版本

# 配置参数
CONDA_ENV="ragen"
PYTHON_CMD="python train.py --config-name aiops-grpo-sft"
LOG_DIR="./logs"
MAX_RETRIES=1000  # 最大重试次数，设置为-1表示无限重试
RESTART_DELAY=30  # 重启前等待时间（秒）
OOM_DELAY=120     # OOM后额外等待时间（秒）

# CUDA设备限制 - 只允许使用设备2和3
export CUDA_VISIBLE_DEVICES="2,3"

# 创建日志目录
mkdir -p "$LOG_DIR"

# 日志文件
MAIN_LOG="$LOG_DIR/training_$(date +%Y%m%d_%H%M%S).log"
ERROR_LOG="$LOG_DIR/errors_$(date +%Y%m%d_%H%M%S).log"
RESTART_LOG="$LOG_DIR/restarts_$(date +%Y%m%d_%H%M%S).log"

# 初始化计数器
attempt=1
total_crashes=0
oom_crashes=0
other_crashes=0

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${GREEN}[INFO]${NC} $timestamp - $message" | tee -a "$MAIN_LOG"
}

log_warning() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${YELLOW}[WARN]${NC} $timestamp - $message" | tee -a "$MAIN_LOG"
}

log_error() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${RED}[ERROR]${NC} $timestamp - $message" | tee -a "$MAIN_LOG" "$ERROR_LOG"
}

# 信号处理函数（简化版本）
handle_signal() {
    log_info "接收到退出信号，正在停止..."
    exit 0
}

# 设置信号处理
trap handle_signal SIGINT SIGTERM

# 检查conda环境
check_conda_env() {
    if ! conda info --envs | grep -q "$CONDA_ENV"; then
        log_error "Conda环境 '$CONDA_ENV' 不存在！"
        exit 1
    fi
}

# 分析崩溃原因
analyze_crash() {
    local exit_code=$1
    local log_content="$2"
    
    # 检查Ray Actor死亡情况
    if echo "$log_content" | grep -qi "ActorDiedError\|ray actor\|worker process has died\|actor died"; then
        return 4  # Ray Actor死亡
    fi
    
    # 检查OOM情况
    if echo "$log_content" | grep -qi "out of memory\|cuda out of memory\|oom\|killed"; then
        return 1  # OOM
    fi
    
    # 检查其他GPU错误
    if echo "$log_content" | grep -qi "cuda error\|gpu\|device"; then
        return 2  # GPU错误
    fi
    
    # 检查网络/IO错误
    if echo "$log_content" | grep -qi "connection\|network\|io error"; then
        return 3  # 网络/IO错误
    fi
    
    return 0  # 其他错误
}

# 主函数
main() {
    log_info "=== PPO训练自动重启脚本启动（无清理版本）==="
    log_info "Conda环境: $CONDA_ENV"
    log_info "训练命令: $PYTHON_CMD"
    log_info "最大重试次数: $MAX_RETRIES"
    log_info "CUDA可见设备: $CUDA_VISIBLE_DEVICES"
    
    # 检查conda环境
    check_conda_env
    
    while true; do
        # 检查是否达到最大重试次数
        if [[ $MAX_RETRIES -gt 0 && $attempt -gt $MAX_RETRIES ]]; then
            log_error "已达到最大重试次数 ($MAX_RETRIES)，退出"
            break
        fi
        
        log_info "=== 第 $attempt 次尝试开始 ==="
        
        # 记录开始时间
        start_time=$(date +%s)
        
        # 激活conda环境并执行训练
        {
            eval "$(conda shell.bash hook)"
            conda activate "$CONDA_ENV"
            
            # 执行训练命令
            $PYTHON_CMD
        } 2>&1 | tee "${LOG_DIR}/attempt_${attempt}_$(date +%Y%m%d_%H%M%S).log"
        
        # 获取退出状态
        exit_status=${PIPESTATUS[0]}
        end_time=$(date +%s)
        duration=$((end_time - start_time))
        
        if [[ $exit_status -eq 0 ]]; then
            log_info "训练成功完成！总用时: ${duration}秒"
            break
        else
            total_crashes=$((total_crashes + 1))
            
            # 获取最后的日志内容用于分析
            last_log="${LOG_DIR}/attempt_${attempt}_$(date +%Y%m%d_%H%M%S).log"
            if [[ -f "$last_log" ]]; then
                log_content=$(tail -100 "$last_log")
            else
                log_content=""
            fi
            
            # 分析崩溃原因
            analyze_crash $exit_status "$log_content"
            crash_type=$?
            
            case $crash_type in
                1)
                    oom_crashes=$((oom_crashes + 1))
                    log_error "检测到OOM错误 (第${oom_crashes}次)，退出状态: $exit_status，运行时长: ${duration}秒"
                    sleep_time=$((RESTART_DELAY + OOM_DELAY))
                    ;;
                2)
                    log_error "检测到GPU错误，退出状态: $exit_status，运行时长: ${duration}秒"
                    sleep_time=$((RESTART_DELAY + 60))
                    ;;
                3)
                    log_error "检测到网络/IO错误，退出状态: $exit_status，运行时长: ${duration}秒"
                    sleep_time=$RESTART_DELAY
                    ;;
                4)
                    log_error "检测到Ray Actor死亡错误，退出状态: $exit_status，运行时长: ${duration}秒"
                    sleep_time=$RESTART_DELAY
                    ;;
                *)
                    other_crashes=$((other_crashes + 1))
                    log_error "检测到其他错误，退出状态: $exit_status，运行时长: ${duration}秒"
                    sleep_time=$RESTART_DELAY
                    ;;
            esac
            
            # 记录重启信息
            echo "$(date '+%Y-%m-%d %H:%M:%S') - 尝试 $attempt 失败，退出状态: $exit_status，运行时长: ${duration}秒，崩溃类型: $crash_type" >> "$RESTART_LOG"
            
            log_warning "等待 $sleep_time 秒后重启... (总崩溃: $total_crashes, OOM: $oom_crashes, 其他: $other_crashes)"
            sleep "$sleep_time"
        fi
        
        attempt=$((attempt + 1))
    done
    
    log_info "=== 训练脚本结束 ==="
    log_info "总尝试次数: $((attempt-1))"
    log_info "总崩溃次数: $total_crashes"
    log_info "OOM崩溃次数: $oom_crashes"
    log_info "其他崩溃次数: $other_crashes"
}

# 运行主函数
main "$@"